/*
 * 文件名：ReconnectingThriftClient.java
 * 描述： ReconnectingThriftClient.java
 * 修改人：Lo&Wang
 * 修改时间：2016年11月24日
 * 修改内容：新增
 */
package com.lowang.test.thrift;

import java.lang.reflect.InvocationHandler;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.lang.reflect.Proxy;
import java.util.Set;

import org.apache.thrift.TServiceClient;
import org.apache.thrift.transport.TTransport;
import org.apache.thrift.transport.TTransportException;

import com.google.common.collect.Sets;

/**
 * @author zhengmo
 */
public class ReconnectingThriftClient {
    private static final Set<Integer> RESTARTABLE_CAUSES = Sets.newHashSet(TTransportException.NOT_OPEN, TTransportException.END_OF_FILE, TTransportException.TIMED_OUT,
            TTransportException.UNKNOWN);
    public static final class Options {
        private int numRetries;
        private long timeBetweenRetries;
        public Options(int numRetries, long timeBetweenRetries) {
            this.numRetries = numRetries;
            this.timeBetweenRetries = timeBetweenRetries;
        }
        /**
         * 设置numRetries.
         * @return 返回numRetries
         */
        public int getNumRetries() {
            return numRetries;
        }
        /**
         * 获取numRetries.
         * @param numRetries 要设置的numRetries
         */
        public void setNumRetries(int numRetries) {
            this.numRetries = numRetries;
        }
        /**
         * 设置timeBetweenRetries.
         * @return 返回timeBetweenRetries
         */
        public long getTimeBetweenRetries() {
            return timeBetweenRetries;
        }
        /**
         * 获取timeBetweenRetries.
         * @param timeBetweenRetries 要设置的timeBetweenRetries
         */
        public void setTimeBetweenRetries(long timeBetweenRetries) {
            this.timeBetweenRetries = timeBetweenRetries;
        }
        public static Options defaults() {
            return new Options(5, 10000L);
        }
    }
    @SuppressWarnings("unchecked")
    public static <T extends TServiceClient, C> C wrap(T baseClient, Class<?> clientInterface, Options options) {
        Object proxyObject = Proxy.newProxyInstance(clientInterface.getClassLoader(), new Class<?>[]{clientInterface},
                new ReconnectionProxy<T>(baseClient, options.numRetries, options.timeBetweenRetries));
        return (C) proxyObject;
    }
    private static class ReconnectionProxy<T extends TServiceClient> implements InvocationHandler {
        private final T baseClient;
        private final int maxRetires;
        private final long timeBetweenRetries;
        @Override
        public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
            try {
                return method.invoke(baseClient, args);
            } catch (InvocationTargetException e) {
                if (e.getTargetException() instanceof TTransportException) {
                    TTransportException ex = (TTransportException) e.getTargetException();
                    if (RESTARTABLE_CAUSES.contains(ex.getType())) { //need reconnecting
                        reconnectOrThrowException(baseClient.getInputProtocol().getTransport(), maxRetires, timeBetweenRetries);
                        return method.invoke(proxy, args);
                    }
                }
                throw e;
            }
        }
        public ReconnectionProxy(T baseClient, int maxRetires, long timeBetweenRetries) {
            this.baseClient = baseClient;
            this.maxRetires = maxRetires;
            this.timeBetweenRetries = timeBetweenRetries;
        }
    }
    private static void reconnectOrThrowException(TTransport transport, int maxRetries, long timeBetweenRetries) throws TTransportException {
        int errors = 0;
        transport.close();
        System.out.println("begin connect...");
        while (errors < maxRetries) {
            try {
                transport.open();
            } catch (TTransportException e) {
                errors++;
                if (errors < maxRetries) {
                    try {
                        Thread.sleep(timeBetweenRetries);
                    } catch (InterruptedException e1) {
                        throw new RuntimeException(e);
                    }
                }
            }
        }
        if (errors >= maxRetries) {
            throw new TTransportException("Failed to reconnect");
          }
    }
}
